# 这里是条件GAN的代码。# 这里要多传入一个labels,这个样子就可以输出你想要的数字。
import time
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import utils, datasets, transforms
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from IPython.display import HTML
import cganclass_function as mycf
import time

# 这里算是主函数吧。
learning_rate, epoch, batch_size, nc, device, num_work = mycf.hyperparameters()  # 返回一些超参数
data = mycf.get_data(batch_size)
print(data)

# 网络的实例化并且赋予其初始权重。
# G = mycf.generator().to(device)
# D = mycf.discriminator().to(device)
# G.apply(mycf.weight_init)
# D.apply(mycf.weight_init)
G = torch.load('./G_CGAN.pth')
D = torch.load('./D_CGAN.pth')
# 这里开始定义一些基本的东西
loss_function = nn.BCELoss()
optimizer_G = torch.optim.Adam(G.parameters(), lr=learning_rate, betas=(0.5, 0.999))
optimizer_D = torch.optim.Adam(D.parameters(), lr=learning_rate, betas=(0.5, 0.999))

# ---------------------------------------分割线-------------------------------------------------
# label one_hot for G
label_1hots = torch.zeros(10, 10)
print(label_1hots.shape)
for i in range(10):
    label_1hots[i, i] = 1
label_1hots = label_1hots.view(10, 10, 1, 1).to(device)  # 这里将labels转化成了一个10通道的向量，前面那个10表示的应该是分类的数量
print(label_1hots[1].shape)  # 在这里创造一个10维的one_hot编码，到了训练的时候就可以直接锁定这个对应的编码。

# Label one-hot for D
label_fills = torch.zeros(10, 10, 32, 32)  # 10个分类加上10个通道数
print(label_fills[1][1].shape)
ones = torch.ones(32, 32)  # 全部都是等于1
for i in range(10):
    label_fills[i][i] = ones

label_fills = label_fills.to(device)
fixed_label = label_1hots[torch.arange(10).repeat(10).sort().values]
print(fixed_label.shape)  # 在这里由于要带上标签，所以输入的时候100表示的是噪声，10表示的是分类。
# 这里要对上面的进行解释arrange函数指的是生成一个0到9的张量，repeat表示的是重复几次，sort是进行排序，value指的是将数值取出来，因为排序还会返回一个索引，这边是不需要索引的。
# -------------------------------------分割线-----------------------------------------------------------

# 接下来开始进行训练
t1 = float('inf')  # 在这里创造一个inf
t2 = float('inf')  # 在这里创造一个inf

for i in range(epoch):
    begin = time.time()
    sum_loss_D = 0  # 这里进行初始化，用于计算loss值。
    sum_loss_G = 0
    cnt = 1
    for images, labels in data:
        cnt = cnt + 1  # 这里是计数器加1
        images = images.to(device)
        labels = labels.to(device)
        size = images.size(0)  # 这里是取出batch_size的大小。

        G_label = label_1hots[labels]
        D_label = label_fills[labels]
        # print(G_label)
        # print(G_label.shape)
        # start to train D!
        D.zero_grad()  # 将梯度清零处理。
        output = D(images, D_label).view(-1)  # 在这里将数据传进去。输出是一个4维张量，所以要进行view的操作。
        real_loss = loss_function(output, torch.ones_like(output, device=device))  # 在这里要和1进行对比。
        # print(real_loss)
        real_loss.backward()

        # 接下来是使用判别器来判别一下生成器生成的图像。
        noise = torch.randn(size, 100, 1, 1).to(device)  # 这里的噪声传入是一个100的通道，也就是一个100维的向量。
        fake_image = G(noise, G_label)  # 在这里可以导出一张假的图像。
        # print(fake_image.shape)
        output = D(fake_image.detach(), D_label).view(-1)  # 这里训练的是D,但是fake这张图像是要经过G的，训练D的时候不需要训练G所以要加上一个detach.
        fake_loss = loss_function(output, torch.zeros_like(output, device=device))  # 这里和0进行对比。
        fake_loss.backward()

        D_loss = fake_loss + real_loss
        sum_loss_D += D_loss.item()  # 这里求出每次的损失值

        optimizer_D.step()

        # 接下来开始训练生成器。
        G.zero_grad()
        output = D(fake_image, D_label).view(-1)  # 在这里由于开始训练G所以不需要在fake_image上面加上detach，这个样子就可以优化G.
        G_loss = loss_function(output,
                               torch.ones_like(output, device=device))  # 在这里要和1进行对比，因为我希望G这个可以迷惑D,也就是说假图像可以迷惑住D.s
        G_loss.backward()
        sum_loss_G += G_loss.item()
        optimizer_G.step()

    # if (D_loss.item() < t1) & (G_loss.item() < t2):
    #     t1 = D_loss.item()
    #     t2 = G_loss.item()
    torch.save(G, './G_CGAN_new.pth')
    torch.save(D, './D_CGAN_new.pth')

    end = time.time()
    run_time = round(end - begin, 5)
    print('Time:{}, Epoch:{}, D_loss:{}, G_loss:{}'.format(run_time, i + 1, round(sum_loss_D / cnt),
                                                           round(sum_loss_G / cnt)))
